﻿using LxBasic;

namespace LxLESS
{
    public partial class MGZone : IBuildMGZone
    {
        //这部分字段不是通过方法设置的，而是直接设置
        public MGCell[] InnCells;
        public MGInnFace[] InnFaces;
        public MGBouFace[] BouFaces;

        public IndexRange[] Bounds;
        public IndexRange InnerCellIdx => new IndexRange(InnCells.Length);
        public IndexRange InnerFaceIdx => new IndexRange(InnFaces.Length);
        #region IProvideZoneTopology
        (int OSideCell, int NSideCell) IBuildMGZone.GetInnFaceAdjCell(int faceID)
        {
            ref var f = ref InnFaces[faceID];
            return (f.OSideCell, f.NSideCell);
        }
        IEnumerable<int> IBuildMGZone.GetInnCellAdjInnCelll(int cellId)
        {
            var cell = InnCells[cellId];
            for (int i1 = cell.InnOFaceEnd; i1 < cell.InnNFaceEnd; i1++)
            {
                yield return InnFaces[cell.AdjFaces[i1]].OSideCell;
            }
            for (int i1 = 0; i1 < cell.InnOFaceEnd; i1++)
            {
                yield return InnFaces[cell.AdjFaces[i1]].NSideCell;
            }
        }
        int IBuildMGZone.GetBoundsCount() => Bounds.Length;

        int IBuildMGZone.GetBouFaceOSideCell(int bouFaceId)
        {
            return BouFaces[bouFaceId].OSideCell;
        }
        (IndexRange CellIdx, IndexRange FaceIdx) IBuildMGZone.GetInnerIdx()
        {
            return (new IndexRange(InnCells.Length), new IndexRange(InnFaces.Length));
        }
        IndexRange IBuildMGZone.GetBoundIdx(int boundID)
        {
            return Bounds[boundID];
        }
        #endregion
    }
    public partial class MGZone
    {
        public LESSZone InLESSZone;
        public double[] e_Common;

        public double[] e;
        public double[] InnAO;
        public double[] InnAN;
        public double[] BouAH;

        #region 外部调用
        public void CreateArray(LESSZone inLESSZone, double[] eCommon)
        {
            CheckCreator.ArrayStrict(ref e, InnerCellIdx.GetCount());
            CheckCreator.ArrayStrict(ref InnAO, InnFaces.Length);
            CheckCreator.ArrayStrict(ref InnAN, InnFaces.Length);
            CheckCreator.ArrayStrict(ref BouAH, BouFaces.Length);
            this.e_Common = eCommon;
            this.InLESSZone = inLESSZone;
        }

        #endregion
        protected unsafe double calALUProductValue(ref MGCell cell, double[] fromArray)
        {
            double aluprotuct = 0;
            fixed (int* fix = cell.AdjFaces)
            {
                int* f = fix, endp1 = fix + cell.InnOFaceEnd;
                for (; f < endp1; ++f)
                {
                    aluprotuct += InnAO[*f] * fromArray[InnFaces[*f].NSideCell];
                }
                endp1 = fix + cell.InnNFaceEnd;
                for (; f < endp1; ++f)
                {
                    aluprotuct += InnAN[*f] * fromArray[InnFaces[*f].OSideCell];
                }
                endp1 = fix + cell.AdjFaces.Length;
                for (; f < endp1; ++f)
                {
                    aluprotuct += BouAH[*f];
                }
            }
            return aluprotuct;
        }
    }
    public partial class MGZone : IZoneIteration
    {
        public void ExplicitIterate_Step1(ThreadInfo threadInfo)
        {
            (int i1, int endp1) = InnerCellIdx.Split(threadInfo);
            Array.Copy(e, i1, e_Common, i1, endp1 - i1);
        }
        public void ExplicitIterate_Step2(ThreadInfo threadInfo)
        {
            jacobiCal(e_Common, e, threadInfo);
        }
        public void JacobiIterate_Step1(ThreadInfo threadInfo)
        {
            jacobiCal(e, e_Common, threadInfo);
        }
        public void JacobiIterate_Step2(ThreadInfo threadInfo)
        {
            jacobiCal(e_Common, e, threadInfo);
        }
        protected void jacobiCal(double[] fromArray, double[] toArray, ThreadInfo threadInfo)
        {
            (int i1, int endp1) = InnerCellIdx.Split(threadInfo);
            for (; i1 < endp1; i1++)
            {
                ref MGCell cell = ref InnCells[i1];
                toArray[i1] = (cell.r - calALUProductValue(ref cell, fromArray)) / cell.ADiag;
            }
        }
        public void AMGIte_Restrict_Are(ThreadInfo threadInfo)
        {
            MGZone coarse = GetCoarserZone();
            var coarseCells = coarse.InnCells;
            var coarsee = coarse.e;
            (int i1, int endp1) = coarse.InnerCellIdx.Split(threadInfo);
            for (; i1 < endp1; i1++)
            {
                ref MGCell coarCell = ref coarseCells[i1];
                coarCell.ADiag = restrict_Ar(ref coarCell);
                coarsee[i1] = coarCell.r / coarCell.ADiag;
            }
            //A_LU
            coarse.HelpRestrict_ALU(InnAO, InnAN, threadInfo);
        }
        public void AMGIte_Restrict_re(ThreadInfo threadInfo)
        {
            MGZone coarse = GetCoarserZone();
            var coarseCells = coarse.InnCells;
            var coarsee = coarse.e;
            (int i1, int endp1) = coarse.InnerCellIdx.Split(threadInfo);
            for (; i1 < endp1; i1++)
            {
                ref MGCell coarCell = ref coarseCells[i1];
                restrict_r(ref coarCell);
                coarsee[i1] = coarCell.r / coarCell.ADiag;// coarseAd[i1];
            }
        }
        private double restrict_Ar(ref MGCell coarseCell)
        {
            double r = 0, atemp = 0;
            foreach (int cell in coarseCell.FineCells)
            {
                ref MGCell fineCell = ref InnCells[cell];
                atemp += fineCell.ADiag;
                r += fineCell.r - calALUProductValue(ref fineCell, e) - fineCell.ADiag * e[cell];
            }
            foreach (int face in coarseCell.InnerFineFaces)
            {
                atemp += InnAO[face] + InnAN[face];
            }
            coarseCell.r = r;
            return atemp;
        }
        private void restrict_r(ref MGCell coarseCell)
        {
            double r = 0;
            foreach (int cell in coarseCell.FineCells)
            {
                ref MGCell fineCell = ref InnCells[cell];
                r += fineCell.r - calALUProductValue(ref fineCell, e) - fineCell.ADiag * e[cell];
            }
            coarseCell.r = r;
        }
        public void AMGIte_ReceiveProlong(ThreadInfo threadInfo)
        {
            GetCoarserZone().HelpProlong(e, threadInfo);
        }
        public void HelpRestrict_ALU(double[] fineALUO, double[] fineALUN, ThreadInfo threadInfo)
        {
            (int i1, int endp1) = InnerFaceIdx.Split(threadInfo);
            //A_LU
            if (InLESSZone.IsSymmetric)
            {
                for (; i1 < endp1; i1++)
                {
                    ref MGInnFace coarseFace = ref InnFaces[i1];
                    double aluo = 0;
                    foreach (int fineface in coarseFace.FineFaces_O2N)
                    {
                        aluo += fineALUO[fineface];
                    }
                    foreach (int fineface in coarseFace.FineFaces_N2O)
                    {
                        aluo += fineALUN[fineface];
                    }
                    InnAO[i1] = InnAN[i1] = aluo;
                }
            }
            else
            {
                for (; i1 < endp1; i1++)
                {
                    ref MGInnFace coarseFace = ref InnFaces[i1];
                    double aluo = 0, alun = 0;
                    foreach (int fineface in coarseFace.FineFaces_O2N)
                    {
                        aluo += fineALUO[fineface];
                        alun += fineALUN[fineface];
                    }
                    foreach (int fineface in coarseFace.FineFaces_N2O)
                    {
                        aluo += fineALUN[fineface];
                        alun += fineALUO[fineface];
                    }
                    InnAO[i1] = aluo;
                    InnAN[i1] = alun;
                }
            }
            (i1, endp1) = new IndexRange(BouAH.Length).Split(threadInfo);
            Array.Clear(BouAH, i1, endp1 - i1);
        }
        public void HelpProlong(double[] e_FinerZone, ThreadInfo threadInfo)
        {
            (int i1, int endp1) = InnerCellIdx.Split(threadInfo);
            for (; i1 < endp1; i1++)
            {
                double e = this.e[i1];
                foreach (int finecell in InnCells[i1].FineCells)
                {
                    e_FinerZone[finecell] += e;
                }
            }
        }
        public MGZone GetCoarserZone()
        {
            return InLESSZone.MGZones[AMGGlobal.SolvingLevel + 1];
        }
        public LESSZone GetLessZone()
        {
            return InLESSZone;
        }

        public void CalResidual(ThreadInfo threadInfo)
        {
            (int i1, int endp1) = InnerCellIdx.Split(threadInfo);
            RMSResidualTool res = new RMSResidualTool();
            for (; i1 < endp1; i1++)
            {
                ref MGCell cell = ref InnCells[i1];
                double future_x = (cell.r - calALUProductValue(ref cell, e)) / cell.ADiag;
                res.Include(e[i1], future_x);
            }
            lock (MultiThreads.GlobalLock)
            {
                InLESSZone.Residual.Include(res);
            }
        }
        public (bool isValid, double res) GetResidualReset()
        {
            double res = InLESSZone.Residual.GetRMS(InnCells.Length);
            InLESSZone.Residual = new RMSResidualTool();
            return (true, res);
        }
        public void BiCGStabIte_Ini_r_rho1(ThreadInfo threadInfo)
        {
            var r = GetBiCGStabSolver().r;
            (int i1, int endp1) = InnerCellIdx.Split(threadInfo);
            Float128 temp = new();
            for (; i1 < endp1; i1++)
            {
                ref MGCell cell = ref InnCells[i1];
                double rt = r[i1] = cell.r - calALUProductValue(ref cell, e) - cell.ADiag * e[i1];
                temp += rt * rt;
            }
            ThreadsMath.MemoryAdd(temp);
        }
        public void BiCGStabIte_Ini_r0_pArray_yArray(ThreadInfo threadInfo)
        {
            (int i1, int endp1) = InnerCellIdx.Split(threadInfo);
            BiCGStabZone bicg = GetBiCGStabSolver();
            double[] r0 = bicg.r0, p = bicg.p, r = bicg.r, y = bicg.y;
            for (; i1 < endp1; i1++)
            {
                double pt = r0[i1] = p[i1] = r[i1];
                y[i1] = pt / InnCells[i1].ADiag;
            }
        }
        public void BiCGStabIte_vArray_alpha1(ThreadInfo threadInfo)
        {
            BiCGStabZone bicg = GetBiCGStabSolver();
            double[] v = bicg.v, y = bicg.y, r0 = bicg.r0;
            (int i1, int endp1) = InnerCellIdx.Split(threadInfo);
            Float128 temp = new();
            for (; i1 < endp1; i1++)
            {
                ref MGCell cell = ref InnCells[i1];
                double vt = v[i1] = cell.ADiag * y[i1] + calALUProductValue(ref InnCells[i1], y);
                temp += (r0[i1] * vt);
            }
            ThreadsMath.MemoryAdd(temp);
        }
        public void BiCGStabIte_sArray_zArray(ThreadInfo threadInfo)
        {
            double alpha = BiCGStabGlobal.alpha;
            BiCGStabZone bicg = GetBiCGStabSolver();
            double[] r = bicg.r, v = bicg.v, z = bicg.z;
            (int i1, int endp1) = InnerCellIdx.Split(threadInfo);
            for (; i1 < endp1; i1++)
            {
                ref double rref = ref r[i1];
                rref -= alpha * v[i1];
                z[i1] = rref / InnCells[i1].ADiag;
            }
        }
        public void BiCGStabIte_tArray_omega1(ThreadInfo threadInfo)
        {
            BiCGStabZone bicg = GetBiCGStabSolver();
            double[] s = bicg.r, z = bicg.z, t = bicg.t;
            (int i1, int endp1) = InnerCellIdx.Split(threadInfo);
            Float128 temp1 = new(), temp2 = new();
            for (; i1 < endp1; i1++)
            {
                ref MGCell cell = ref InnCells[i1];
                double tt = t[i1] = cell.ADiag * z[i1] + calALUProductValue(ref InnCells[i1], z);
                temp1 += tt * s[i1];
                temp2 += tt * tt;
            }
            ThreadsMath.MemoryAdd(temp1, temp2);
        }
        public void BiCGStabIte_xArray(ThreadInfo threadInfo)
        {
            GetBiCGStabSolver().BiCGStab_Cal_xArray(InnerCellIdx, e, threadInfo);
        }
        public void BiCGStabIte_rArray_rho1(ThreadInfo threadInfo)
        {
            GetBiCGStabSolver().BiCGStab_rArray_rho1(InnerCellIdx, threadInfo);
        }
        public void BiCGStabIte_pArray_yArray(ThreadInfo threadInfo)
        {
            double beta = BiCGStabGlobal.beta; double omega = BiCGStabGlobal.omega;
            BiCGStabZone bicg = GetBiCGStabSolver();
            double[] p = bicg.p, r = bicg.r, v = bicg.v, y = bicg.y;
            (int i1, int endp1) = InnerCellIdx.Split(threadInfo);
            for (; i1 < endp1; i1++)
            {
                ref double pt = ref p[i1];
                pt = r[i1] + beta * (pt - omega * v[i1]);
                y[i1] = pt / InnCells[i1].ADiag;
            }
        }
        public BiCGStabZone GetBiCGStabSolver()
        {
            return InLESSZone.BiCGZone;
        }

        public bool IsSymmetric()
        {
            return InLESSZone.IsSymmetric;
        }
    }
}